# -*- coding:UTF-8 -*-
import ruamel.yaml
from ruamel import yaml

class YamlUtil(object):
    def __init__(self, path, param=None):
        self.path = path  # 文件路径
        # print(self.path)
        self.param = param  # 读取信息的key,不传默认获取所有数据


    # 获取yaml文件中的数据
    def get_data(self):
        with open(self.path, 'rb') as y:
            cont = y.read()  # 获取yaml文件中的所有信息
        yaml.warnings({'YAMLLoadWarning': False})  # 禁用加载器warnings报警
        cf = yaml.load(cont)  # 将bytes格式转成dict格式
        y.close()  # 关闭文件
        if self.param == None:
            print(cf)
            return cf  # 返回所有数据
        else:
            return cf.get(self.param)


    def get_ruamel(self):
        with open(self.path, "r", encoding="utf-8") as docs:
            try:
                alldata = ruamel.yaml.safe_load(docs)
            except ruamel.yaml.YAMLError as exc:
                print(exc)
        docs.close()
        # print(alldata)
        return alldata


    def get_case(self):
        try:
            with open(self.path, "r", encoding="utf-8") as docs:
                try:
                    alldata = ruamel.yaml.safe_load(docs)
                except ruamel.yaml.YAMLError as exc:
                    print(exc)
            docs.close()
            # print(alldata)
            return alldata['TestCases']
        except Exception as e:
            print("读取测试用例异常！！-------%s" % e)

    def get_conf(self):
        with open(self.path, "r", encoding="utf-8") as docs:
            try:
                alldata = ruamel.yaml.safe_load(docs)
            except ruamel.yaml.YAMLError as exc:
                print(exc)
        docs.close()
        # print(alldata)
        return alldata["TEST_SETTING"]

    def getData(self):
        with open(self.path, "r", encoding="utf-8") as docs:
            try:
                alldata = ruamel.yaml.safe_load(docs)
            except ruamel.yaml.YAMLError as exc:
                print(exc)
        docs.close()
        # print(alldata)
        return alldata


    # 修改yaml
    def set_ruamel(self,i, setKey, setValue):
        # # 打开
        # with open(self.path, "r", encoding="utf-8") as docs:
        #     try:
        #         alldata = ruamel.yaml.safe_load(docs)
        #     except ruamel.yaml.YAMLError as exc:
        #         print(exc)
        allData = YamlUtil(self.path).get_ruamel()
        # 修改
        allData['TestCases'][i][setKey] = setValue
        # 写入
        with open(self.path, 'w', encoding='utf8') as outfile:
            ruamel.yaml.round_trip_dump(allData, outfile, default_flow_style=False, allow_unicode=True)
        outfile.close()
